{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "#sys.path.append('/home/luca/GitRepositories/Brancher')\n",
    "sys.path.append('/home/luca/GitRepositories/Brancher')\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will explain how to integrate deep learning models into your Brancher pipline. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bulding models with NN components\n",
    "\n",
    "All pytorch functions can be used in Brancher. These functions need to be imported from the Brancher function module which converts all PyTorch functions acting on torch.Tensor to Brancher functions acting on brancher.Variable.\n",
    "\n",
    "Let's get started by defining a stochastic convolutional network on the MNIST dataset. The first step is to import the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "\n",
    "# Data\n",
    "image_size = 28\n",
    "num_classes = 10\n",
    "train = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=None)\n",
    "test = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=None)\n",
    "dataset_size = len(train)\n",
    "input_variable = np.reshape(train.train_data.numpy(), newshape=(dataset_size, 1, image_size, image_size))\n",
    "output_labels = train.train_labels.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a Brancher model, datasets are stored in empirical variables. These are random variables that sample minibatches from a dataset.\n",
    "In a supervised problem we need two empirical variables, one for the input images and the other for the labels. However, these two variables need to be sampled jointly as each image should be associated to a particular label. In Brancher, we can implement by this coupling by greating an additional RandomIndices variable that is a parent of both images and labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brancher.standard_variables import EmpiricalVariable as Empirical\n",
    "from brancher.standard_variables import RandomIndices\n",
    "\n",
    "# Data sampling model\n",
    "minibatch_size = 7\n",
    "\n",
    "minibatch_indices = RandomIndices(dataset_size=dataset_size, batch_size=minibatch_size, \n",
    "                                  name=\"indices\", is_observed=True)\n",
    "\n",
    "x = Empirical(input_variable, indices=minibatch_indices, \n",
    "              name=\"x\", is_observed=True)\n",
    "\n",
    "labels = Empirical(output_labels, indices=minibatch_indices, \n",
    "                   name=\"labels\", is_observed=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, we can import the PyTorch conv2d function from brancher.functions to build a stochastic 2D convolutional layer with gaussian weights. We can do this by defining oth weights and the input image as Randomvariables and using the Brancher conv2d function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brancher import functions as BF\n",
    "\n",
    "from brancher.standard_variables import NormalVariable as Normal\n",
    "from brancher.standard_variables import BinomialVariable as Binomial\n",
    "\n",
    "in_channels = 1\n",
    "out_channels = 5\n",
    "image_size = 28\n",
    "#x = Normal(loc=np.zeros((in_channels, image_size, image_size)),\n",
    "#           scale=1.,\n",
    "#           name=\"x\")\n",
    "Wk = Normal(loc=np.zeros((out_channels, in_channels, 3, 3)),\n",
    "            scale=np.ones((out_channels, in_channels, 3, 3)),\n",
    "            name=\"Wk\")\n",
    "z = Normal(BF.conv2d(x, Wk, padding=1), 0.1, name=\"z\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here the random input a is convolved with the random filters W. We can now run the forward pass by sampling from the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(7, 5, 28, 28)\n",
      "6\n"
     ]
    }
   ],
   "source": [
    "print(z.get_sample(num_samples)[\"z\"][0].shape)\n",
    "print(len(z.get_sample(num_samples)[\"z\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that in each of these samples both the input and the weigths are sampled independently.\n",
    "\n",
    "We can now add a linear layer to this result to get a shallow convolutional classifier. We can do this by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 10\n",
    "Wl = Normal(loc=np.zeros((num_classes, image_size*image_size*out_channels)),\n",
    "            scale=np.ones((num_classes, image_size*image_size*out_channels)),\n",
    "            name=\"Wl\")\n",
    "b = Normal(loc=np.zeros((num_classes, 1)),\n",
    "           scale=np.ones((num_classes, 1)),\n",
    "           name=\"b\")\n",
    "reshaped_z = BF.reshape(z, shape=(image_size*image_size*out_channels, 1))\n",
    "k = Binomial(1, \n",
    "             logits=BF.linear(reshaped_z, Wl, b), \n",
    "             name=\"k\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we needed to reshape the variable z as we need to perform matrix multiplication in the linear layer. Note that in Brancher you never need to explicitly consider the batch dimension. Batch properties are part of the data, not of the model and Brancher handle them automatically!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(7, 10, 1)\n",
      "6\n"
     ]
    }
   ],
   "source": [
    "print(k.get_sample(num_samples)[\"k\"][0].shape)\n",
    "print(len(k.get_sample(num_samples)[\"k\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bulding models using Brancher layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using existing PyTorch models in a Brancher models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
